/*
 * To change this license header, choose License Headers in Project Properties.
 * To change this template file, choose Tools | Templates
 * and open the template in the editor.
 */
package ui;

/**
 *
 * @author Raise
 */
import utils.RecognitionResult;
import utils.ImageProcessor;
import network.CNN;
import date.TestImageManager;
import javax.swing.*;
import java.awt.*;
import java.awt.event.*;
import java.awt.image.BufferedImage;
import java.util.ArrayList;
import java.util.List;
import java.io.File;

public class AutoTestUI extends JFrame{
     private CNN model;
    private TestImageManager imageManager;
    private ImageProcessor imageProcessor;
    private int currentIndex = 0;
    
    // UI组件
    private JLabel imageLabel;
    private JLabel imageNameLabel;
    private JLabel resultLabel;
    private JLabel confidenceLabel;
    private JTextArea detailTextArea;
    private JLabel indexLabel;
    private JButton prevButton;
    private JButton nextButton;
    private JButton recognizeAllButton;
    private JLabel statusLabel;
    
    public AutoTestUI(CNN model) {
        this.model = model;
        this.imageManager = new TestImageManager();
        this.imageProcessor = new ImageProcessor();
        
        initializeUI();
        loadImages();
        showImage(0);
    }
    
    private void initializeUI() {
        setTitle("MNIST手写数字识别测试系统");
        setDefaultCloseOperation(JFrame.EXIT_ON_CLOSE);
        setLayout(new BorderLayout());
        
        // 创建主面板
        JPanel mainPanel = new JPanel(new BorderLayout(10, 10));
        mainPanel.setBorder(BorderFactory.createEmptyBorder(10, 10, 10, 10));
        
        // 左侧图片面板
        JPanel imagePanel = createImagePanel();
        // 右侧结果面板
        JPanel resultPanel = createResultPanel();
        // 底部控制面板
        JPanel controlPanel = createControlPanel();
        // 状态栏
        statusLabel = new JLabel(" 就绪");
        statusLabel.setBorder(BorderFactory.createLoweredBevelBorder());
        
        mainPanel.add(imagePanel, BorderLayout.WEST);
        mainPanel.add(resultPanel, BorderLayout.CENTER);
        mainPanel.add(controlPanel, BorderLayout.SOUTH);
        
        add(mainPanel, BorderLayout.CENTER);
        add(statusLabel, BorderLayout.SOUTH);
        
        pack();
        setLocationRelativeTo(null);
        setVisible(true);
    }
    
    private JPanel createImagePanel() {
        JPanel panel = new JPanel(new BorderLayout());
        panel.setBorder(BorderFactory.createTitledBorder("测试图片"));
        panel.setPreferredSize(new Dimension(320, 400));
        
        imageNameLabel = new JLabel("", JLabel.CENTER);
        imageNameLabel.setFont(new Font("微软雅黑", Font.BOLD, 14));
        
        imageLabel = new JLabel("加载中...", JLabel.CENTER);
        imageLabel.setPreferredSize(new Dimension(280, 280));
        imageLabel.setBorder(BorderFactory.createLineBorder(Color.GRAY));
        
        panel.add(imageNameLabel, BorderLayout.NORTH);
        panel.add(imageLabel, BorderLayout.CENTER);
        
        return panel;
    }
    
    private JPanel createResultPanel() {
        JPanel panel = new JPanel(new BorderLayout());
        panel.setBorder(BorderFactory.createTitledBorder("识别结果"));
        panel.setPreferredSize(new Dimension(300, 400));
        
        // 结果显示区域
        JPanel displayPanel = new JPanel(new GridLayout(3, 1, 5, 5));
        displayPanel.setBorder(BorderFactory.createEmptyBorder(10, 10, 10, 10));
        
        resultLabel = new JLabel("等待识别", JLabel.CENTER);
        resultLabel.setFont(new Font("微软雅黑", Font.BOLD, 24));
        resultLabel.setForeground(Color.BLUE);
        
        confidenceLabel = new JLabel("置信度: -", JLabel.CENTER);
        confidenceLabel.setFont(new Font("微软雅黑", Font.PLAIN, 16));
        
        displayPanel.add(resultLabel);
        displayPanel.add(confidenceLabel);
        
        // 详细概率区域
        detailTextArea = new JTextArea(12, 20);
        detailTextArea.setEditable(false);
        detailTextArea.setFont(new Font("Consolas", Font.PLAIN, 12));
        JScrollPane scrollPane = new JScrollPane(detailTextArea);
        scrollPane.setBorder(BorderFactory.createTitledBorder("概率分布"));
        
        panel.add(displayPanel, BorderLayout.NORTH);
        panel.add(scrollPane, BorderLayout.CENTER);
        
        return panel;
    }
    
    private JPanel createControlPanel() {
        JPanel panel = new JPanel(new FlowLayout());
        panel.setBorder(BorderFactory.createTitledBorder("控制面板"));
        
        prevButton = new JButton("上一张");
        nextButton = new JButton("下一张");
        recognizeAllButton = new JButton("批量识别所有图片");
        indexLabel = new JLabel("0/0");
        
        prevButton.addActionListener(e -> showImage(currentIndex - 1));
        nextButton.addActionListener(e -> showImage(currentIndex + 1));
        recognizeAllButton.addActionListener(e -> recognizeAllImages());
        
        panel.add(prevButton);
        panel.add(nextButton);
        panel.add(recognizeAllButton);
        panel.add(indexLabel);
        
        return panel;
    }
    
    private void loadImages() {
        statusLabel.setText(" 正在加载测试图片...");
        System.out.println("开始加载内置测试图片...");
        imageManager.loadAllImages();
        updateNavigation();
        statusLabel.setText(" 图片加载完成");
    }
    
    private void showImage(int index) {
        int totalImages = imageManager.getImageCount();
        if (totalImages == 0) {
            imageLabel.setText("未找到测试图片");
            statusLabel.setText(" 错误: 未找到测试图片");
            return;
        }
        
        if (index < 0) index = totalImages - 1;
        if (index >= totalImages) index = 0;
        
        currentIndex = index;
        
        BufferedImage image = imageManager.getImage(currentIndex);
        String imagePath = imageManager.getImagePath(currentIndex);
        
        if (image != null) {
            // 显示图片
            BufferedImage displayImage = imageProcessor.resizeForDisplay(image, 280, 280);
            imageLabel.setIcon(new ImageIcon(displayImage));
            imageLabel.setText("");
            imageNameLabel.setText("图片: " + new File(imagePath).getName());
            
            // 自动识别当前图片
            recognizeCurrentImage();
        } else {
            imageLabel.setIcon(null);
            imageLabel.setText("图片加载失败");
            imageNameLabel.setText("图片: " + imagePath);
            clearResults();
        }
        
        updateNavigation();
    }
    
    private void recognizeCurrentImage() {
        BufferedImage image = imageManager.getImage(currentIndex);
        if (image == null) return;
        
        try {
            statusLabel.setText(" 正在识别当前图片...");
            double[] processed = imageProcessor.preprocessImage(image);
            double[] result = model.forward(processed);
            int predictedDigit = argmax(result);
            double confidence = result[predictedDigit];
            
            // 更新界面显示
            resultLabel.setText("识别结果: " + predictedDigit);
            confidenceLabel.setText("置信度: " + String.format("%.2f%%", confidence * 100));
            
            // 显示详细概率
            StringBuilder details = new StringBuilder();
            for (int i = 0; i < result.length; i++) {
                details.append(String.format("数字 %d: %.4f (%.2f%%)%n", 
                    i, result[i], result[i] * 100));
            }
            detailTextArea.setText(details.toString());
            
            statusLabel.setText(" 识别完成: 数字 " + predictedDigit);
            
        } catch (Exception e) {
            resultLabel.setText("识别失败");
            confidenceLabel.setText("错误: " + e.getMessage());
            detailTextArea.setText("识别过程中出现错误:\n" + e.getMessage());
            statusLabel.setText(" 识别失败");
        }
    }
    
    private void recognizeAllImages() {
        int totalImages = imageManager.getImageCount();
    // 使用数组来包装变量，这样在lambda中就可以修改了
    final int[] successCount = {0};  // 改为数组
    List<RecognitionResult> results = new ArrayList<>();
    
    // 创建进度对话框
    JDialog progressDialog = new JDialog(this, "批量识别", true);
    progressDialog.setLayout(new BorderLayout());
    JProgressBar progressBar = new JProgressBar(0, totalImages);
    progressBar.setStringPainted(true);
    JLabel progressLabel = new JLabel("正在识别图片...", JLabel.CENTER);
    
    progressDialog.add(progressLabel, BorderLayout.NORTH);
    progressDialog.add(progressBar, BorderLayout.CENTER);
    progressDialog.setSize(300, 100);
    progressDialog.setLocationRelativeTo(this);
    
    // 在新线程中执行批量识别
    new Thread(() -> {
        for (int i = 0; i < totalImages; i++) {
            if (imageManager.isImageLoaded(i)) {
                try {
                    BufferedImage image = imageManager.getImage(i);
                    double[] processed = imageProcessor.preprocessImage(image);
                    double[] result = model.forward(processed);
                    int predictedDigit = argmax(result);
                    double confidence = result[predictedDigit];
                    
                    results.add(new RecognitionResult(predictedDigit, confidence, result, 
                                                     imageManager.getImagePath(i), true));
                    successCount[0]++;  // 修改数组元素
                } catch (Exception e) {
                    results.add(new RecognitionResult(imageManager.getImagePath(i)));
                }
            }
            
            final int current = i + 1;
            SwingUtilities.invokeLater(() -> {
                progressBar.setValue(current);
                progressLabel.setText(String.format("正在识别图片 (%d/%d)...", current, totalImages));
            });
            
            try { Thread.sleep(100); } catch (InterruptedException e) {} // 稍微延迟以便观察
        }
        
        SwingUtilities.invokeLater(() -> {
            progressDialog.dispose();
            showBatchResults(results, successCount[0], totalImages);  // 使用数组元素
        });
    }).start();
    
    progressDialog.setVisible(true);
    }
    
    private void showBatchResults(List<RecognitionResult> results, int successCount, int totalImages) {
        StringBuilder sb = new StringBuilder();
        sb.append(String.format("批量识别完成: %d/%d 成功\n\n", successCount, totalImages));
        sb.append("详细结果:\n");
        sb.append("----------------------------------------\n");
        
        for (RecognitionResult result : results) {
            if (result.isSuccessful()) {
                sb.append(String.format("图片: %s → 数字: %d (置信度: %.2f%%)\n",
                    result.getFileName(), result.getPredictedDigit(), result.getConfidence() * 100));
            } else {
                sb.append(String.format("图片: %s → 识别失败\n", result.getFileName()));
            }
        }
        
        JTextArea resultArea = new JTextArea(20, 50);
        resultArea.setText(sb.toString());
        resultArea.setEditable(false);
        resultArea.setFont(new Font("Consolas", Font.PLAIN, 12));
        
        JScrollPane scrollPane = new JScrollPane(resultArea);
        JOptionPane.showMessageDialog(this, scrollPane, "批量识别结果", JOptionPane.INFORMATION_MESSAGE);
    }
    
    private void updateNavigation() {
        int totalImages = imageManager.getImageCount();
        indexLabel.setText(String.format("%d/%d", currentIndex + 1, totalImages));
        
        // 更新按钮状态
        prevButton.setEnabled(totalImages > 0);
        nextButton.setEnabled(totalImages > 0);
        recognizeAllButton.setEnabled(totalImages > 0);
    }
    
    private void clearResults() {
        resultLabel.setText("等待识别");
        confidenceLabel.setText("置信度: -");
        detailTextArea.setText("");
    }
    
    private int argmax(double[] array) {
        int maxIndex = 0;
        for (int i = 1; i < array.length; i++) {
            if (array[i] > array[maxIndex]) {
                maxIndex = i;
            }
        }
        return maxIndex;
    }
}
